import numpy as np
import matplotlib.pyplot as plt
import matplotlib

# 设置中文字体和解决负号显示问题
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False
# 线性回归模型
from sklearn.linear_model import LinearRegression
# 构建多项式回归模型
from sklearn.preprocessing import PolynomialFeatures
# 数据集划分(训练集和测试集)
from sklearn.model_selection import train_test_split
# 平方损失函数（均方误差损失函数）
from sklearn.metrics import mean_squared_error

'''
1. 生成数据
2. 划分训练集和测试集（当前是验证集）
3. 定义模型（线性回归模型）
4. 训练模型
5. 预算结果，计算误差
'''

# 1. 生成数据
X = np.linspace(-3, 3, 300).reshape(-1, 1)
y = np.sin(X) + np.random.uniform(low=-0.5, high=0.5, size=300).reshape(-1, 1)
print(X.shape)
print(y.shape)
# 画出散点图，需要三个子图
fig, ax = plt.subplots(1, 3, figsize=(15, 4))
ax[0].scatter(X, y, color='y')
ax[1].scatter(X, y, color='b')
ax[2].scatter(X, y, color='g')
# plt.show()

# 2. 划分训练集和测试集（当前是验证集）
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 3. 定义模型（线性回归模型）
model = LinearRegression()

# 一、欠拟合（直线）
X_train1 = X_train
X_test1 = X_test

# 4. 训练模型
model.fit(X_train1, y_train)
# 打印查看模型参数
print(model.coef_)
print(model.intercept_)

# 5. 预算结果，计算误差
y_pred1 = model.predict(X_test1)
# 损失函数
test_loss1 = mean_squared_error(y_test, y_pred1)
train_loss1 = mean_squared_error(y_train, model.predict(X_train1))
print(y_pred1.shape)

# 画出拟合曲线，并写出训练误差和测试误差
ax[0].plot(X, model.predict(X), label='LinearRegression', color='r')
ax[0].text(-3, 1, '测试误差: %.4f' % test_loss1)
ax[0].text(-3, 1.3, '训练误差: %.4f' % train_loss1)
ax[0].set_title('欠拟合')
# plt.show()

# 二、恰好拟合（5次多项式）
poly5 = PolynomialFeatures(degree=5)
X_train2 = poly5.fit_transform(X_train)
X_test2 = poly5.fit_transform(X_test)
print('X_train2.shape', '=', X_train2.shape)
print('X_test2.shape', '=', X_test2.shape)

# 4. 训练模型
model.fit(X_train2, y_train)

# 打印查看模型参数
# print(model.coef_)
# print(model.intercept_)

# 5. 预算结果，计算误差
y_pred2 = model.predict(X_test2)
# 损失函数
test_loss2 = mean_squared_error(y_test, y_pred2)
train_loss2 = mean_squared_error(y_train, model.predict(X_train2))
print(y_pred2.shape)

# 画出拟合曲线，并写出训练误差和测试误差
ax[1].plot(X, model.predict(poly5.fit_transform(X)), label='LinearRegression', color='r')
ax[1].text(-3, 1, '测试误差: %.4f' % test_loss2)
ax[1].text(-3, 1.3, '训练误差: %.4f' % train_loss2)
ax[1].set_title('恰好拟合')


# 三、过度拟合（20次多项式）
poly20 = PolynomialFeatures(degree=20)
X_train3 = poly20.fit_transform(X_train)
X_test3 = poly20.fit_transform(X_test)
print('X_train3.shape', '=', X_train3.shape)
print('X_test3.shape', '=', X_test3.shape)

# 4. 训练模型
model.fit(X_train3, y_train)

# 打印查看模型参数
# print(model.coef_)
# print(model.intercept_)

# 5. 预算结果，计算误差
y_pred3 = model.predict(X_test3)
# 损失函数
test_loss3 = mean_squared_error(y_test, y_pred3)
train_loss3 = mean_squared_error(y_train, model.predict(X_train3))
print(y_pred3.shape)

# 画出拟合曲线，并写出训练误差和测试误差
ax[2].plot(X, model.predict(poly20.fit_transform(X)), label='LinearRegression', color='r')
ax[2].text(-3, 1, '测试误差: %.4f' % test_loss3)
ax[2].text(-3, 1.3, '训练误差: %.4f' % train_loss3)
ax[2].set_title('过度拟合')
plt.show()
